import numpy as np
import torch as pt
from mpl_toolkits import mplot3d
import matplotlib.pyplot as plt
import sys
import os
from sklearn.preprocessing import StandardScaler

np.random.seed(777)
pt.manual_seed(777)

VER = 'v2.0'
ALPHA = 0.01
N_EPOCHS = 1000

# data
data = np.loadtxt('../../../../ML/logic_regression/data/ex2data1.txt', delimiter=',')
X = data[:, :-1]
Y = data[:, -1:]

# scaler
scaler = StandardScaler()
X = scaler.fit_transform(X)

# plot
spr = 1
spc = 3
spn = 0
plt.figure(figsize=[18, 6])

# show data
spn += 1
ax = plt.subplot(spr, spc, spn)
pos_idx = np.isclose(1., Y.ravel())
neg_idx = np.invert(pos_idx)
plt.scatter(X[pos_idx, 0], X[pos_idx, 1], color='r', label='positive')
plt.scatter(X[neg_idx, 0], X[neg_idx, 1], color='b', label='negative')
Xt = pt.Tensor(X)
Y = pt.Tensor(Y)

# model
model = pt.nn.Sequential(
    pt.nn.Linear(2, 1, bias=True),
    pt.nn.Sigmoid()
)
criterion = pt.nn.BCELoss()
vars = model.parameters()
optim = pt.optim.Adam(params=vars, lr=ALPHA)


def acc(ht, yt):
    return (ht > 0.5).eq(yt > 0.5).double().mean()  # ATTENTION


loss_history = np.zeros(N_EPOCHS)
acc_history = np.zeros(N_EPOCHS)
GROUP = int(np.ceil(N_EPOCHS / 20))
for step in range(N_EPOCHS):
    optim.zero_grad()
    ht = model(Xt)
    loss = criterion(ht, Y)
    loss.backward()
    optim.step()
    loss_history[step] = loss
    accv = acc(ht, Y)
    acc_history[step] = accv
    if step % GROUP == 0:
        print(f'#{step + 1}: loss = {loss}, acc = {accv}')
if step % GROUP != 0:
    print(f'#{step + 1}: loss = {loss}, acc = {accv}')

params = list(model.parameters())

spn += 1
plt.subplot(spr, spc, spn)
plt.plot(loss_history)
plt.title('Loss history')

spn += 1
plt.subplot(spr, spc, spn)
plt.plot(acc_history)
plt.title('Accuracy history')

# show hyper-plane
bias = params[1][0].data.numpy()
theta1 = params[0][0, 0].data.numpy()
theta2 = params[0][0, 1].data.numpy()
xmin = np.min(X[:, 0])
xmax = np.max(X[:, 0])
plt_x = np.array([xmin, xmax])
plt_y = - (plt_x * theta1 + bias) / theta2
ax.plot(plt_x, plt_y, 'r-')

# finally show all plotting
plt.show()
